import pytz
import pandas as pd
from datetime import datetime

def convert_timezone(dt, target_timezone='Asia/Shanghai'):
    """
    将时间转换为指定时区。
    
    参数:
        dt: 需要转换的datetime对象或pandas Timestamp
        target_timezone: 目标时区
        
    返回:
        转换后的datetime对象
    """
    # 处理pandas Timestamp
    if pd.api.types.is_datetime64_any_dtype(pd.Series([dt])):
        if dt.tzinfo is None:
            dt = dt.tz_localize('UTC')
        return dt.tz_convert(target_timezone)
    
    # 处理标准datetime
    if dt.tzinfo is None:
        # 如果时间无时区信息，默认当作UTC处理
        dt = dt.replace(tzinfo=pytz.UTC)
    
    return dt.astimezone(pytz.timezone(target_timezone))

def ensure_timezone(dt, default_timezone='Asia/Shanghai'):
    """
    确保datetime对象具有时区信息，如果没有，则添加默认时区。
    
    参数:
        dt: 需要检查的datetime对象或pandas Timestamp
        default_timezone: 默认时区
        
    返回:
        带有时区信息的datetime对象
    """
    # 处理pandas Timestamp
    if pd.api.types.is_datetime64_any_dtype(pd.Series([dt])):
        if dt.tzinfo is None:
            return dt.tz_localize(default_timezone)
        return dt
    
    # 处理标准datetime
    if dt.tzinfo is None:
        return pytz.timezone(default_timezone).localize(dt)
    return dt

def standardize_to_shanghai(dt):
    """
    将时间标准化为上海时区。
    如果时间没有时区信息，会先假设它是UTC时间。
    
    参数:
        dt: 需要标准化的datetime对象或pandas Timestamp
        
    返回:
        上海时区的datetime对象
    """
    return convert_timezone(dt, 'Asia/Shanghai')

def convert_timestamp_to_datetime(timestamp, unit='s', timezone='Asia/Shanghai'):
    """
    将时间戳转换为带时区的日期时间对象。
    
    参数:
        timestamp: UNIX时间戳
        unit: 时间戳单位，'s'表示秒，'ms'表示毫秒
        timezone: 目标时区
        
    返回:
        带时区的datetime对象
    """
    dt = pd.to_datetime(timestamp, unit=unit)
    if timezone:
        # 先设置为UTC，然后转换到目标时区
        dt = dt.tz_localize('UTC').tz_convert(timezone)
    
    return dt 